#include "naive_pairwise.hpp"

#include "MG.hpp"
#include "parcsr_matop.hpp"
// #define DEBUG

template<typename idx_t, typename setup_t, int dof>
idx_t * naive_pairwise(const par_CSRMatrix<idx_t, setup_t, setup_t, dof> * par_fine_mat,
    par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & R_mat, par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & P_mat)
{
    const MPI_Comm comm = par_fine_mat->comm;
    int my_pid; MPI_Comm_rank(comm, & my_pid);
    int num_procs; MPI_Comm_size(comm, & num_procs);
    constexpr MPI_Datatype mpi_idx_type = sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT;

    const idx_t my_F_beg = par_fine_mat->beg_row;
    const idx_t my_num_F = par_fine_mat->end_row - my_F_beg;
    idx_t my_num_C = 0;

    // 因为是naive，所以只考虑对角部分
    CSRMatrix<idx_t, setup_t, setup_t, 1> S_mat(my_num_F, par_fine_mat->diag.ncols, par_fine_mat->diag.nnz);
    constexpr int e_size = dof*dof;
    auto calc_blk_norm = [] (const setup_t * val) {
        setup_t ret = 0.0;
        for (int f = 0; f < e_size; f++) ret += val[f] * val[f];
        return sqrt(ret);
    };
    
    idx_t * loc_map_f2c = new idx_t [my_num_F];
    coar_wells.clear();// 为这一轮粗化清空
    constexpr idx_t UNAGGREGATED = -1;

    #pragma omp parallel
    {
        #pragma omp for schedule(static)
        for (idx_t i = 0; i <= my_num_F; i++) {
            S_mat.row_ptr[i] = par_fine_mat->diag.row_ptr[i];
        }
        #pragma omp for schedule(static)
        for (idx_t p = 0; p <  S_mat.row_ptr[my_num_F]; p ++) {
            S_mat.col_idx[p] = par_fine_mat->diag.col_idx[p];
            S_mat.vals   [p] = calc_blk_norm(par_fine_mat->diag.vals + p * e_size);
        }
        #pragma omp for schedule(static) nowait
        for (idx_t i = 0; i < my_num_F; i++) {
#define PRES_COL_ORDER
#ifdef PRES_COL_ORDER // 保证列号升序（以便coupled和naive在单进程时严格复现）
            sort_vals_with_indices<idx_t, setup_t>(S_mat.row_ptr[i + 1] - S_mat.row_ptr[i], S_mat.vals    + S_mat.row_ptr[i],
                                                                                            S_mat.col_idx + S_mat.row_ptr[i] );
#else // 快排不稳定排序
            my_qsort1<idx_t, setup_t, idx_t, -1>(S_mat.vals + S_mat.row_ptr[i], S_mat.col_idx + S_mat.row_ptr[i],
                0, S_mat.row_ptr[i + 1] - S_mat.row_ptr[i] - 1);
#endif
        }

        // 初始化
        #pragma omp for schedule(static)
        for (idx_t fi = 0; fi < my_num_F; fi++)
            loc_map_f2c[fi] = UNAGGREGATED;

    }

    auto trans_loc2glb = [my_F_beg](const idx_t loc) {
        return loc + my_F_beg;
    };
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
        if (loc_map_f2c[loc_fi] != UNAGGREGATED) {// 之前已经聚合过了
            continue;
        }
        if (iso_wells.find(trans_loc2glb(loc_fi)) != iso_wells.end()){// 井不与别人聚合，单独成一个粗点
            assert(loc_map_f2c[loc_fi] == UNAGGREGATED);// 之前不能被聚合过
            coar_wells.emplace(my_num_C);
            loc_map_f2c[loc_fi] = my_num_C;
            my_num_C ++;
            continue;
        }

        // 从大到小考察邻居
        for (idx_t p = S_mat.row_ptr[loc_fi]; p < S_mat.row_ptr[loc_fi + 1]; p++) {
            const idx_t loc_fj = S_mat.col_idx[p];
            if (loc_fj != loc_fi && // 该邻居不能是自己
                iso_wells.find(trans_loc2glb(loc_fj)) == iso_wells.end() && // 也不能是井
                loc_map_f2c[loc_fj] == UNAGGREGATED) // 也尚未被聚合
            {
                loc_map_f2c[loc_fi] = my_num_C;
                loc_map_f2c[loc_fj] = my_num_C;
                // printf("%d %d => %d\n", loc_fi, loc_fj, my_num_C);
                my_num_C ++;
                break;
            }
        }
    }

    idx_t * C_rows_partits = new idx_t [num_procs + 1];
    C_rows_partits[0] = 0;
    MPI_Allgather(& my_num_C      , 1, sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT,
                C_rows_partits + 1, 1, sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT, comm);
    for (int p = 0; p < num_procs; p++)
        C_rows_partits[p + 1] += C_rows_partits[p];

    // 构造限制矩阵
    R_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, 1>(comm, C_rows_partits, par_fine_mat->rows_partition);
    P_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, 1>(comm, par_fine_mat->cols_partition, C_rows_partits);
#if 1 // 允许跨进程聚合
    const idx_t _cbeg = C_rows_partits[my_pid], _cend = C_rows_partits[my_pid + 1];
    idx_t * glb_map_f2c = new idx_t [my_num_F];// loc_f => glb_c
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
        if (loc_map_f2c[loc_fi] >= 0)
            glb_map_f2c[loc_fi] = loc_map_f2c[loc_fi] + _cbeg;// 累加全局偏移
        else
            glb_map_f2c[loc_fi] = -1;// 暂时置为-1
    }
    {// 处理尚未聚合的点
        idx_t * send_buf = new idx_t [par_fine_mat->commpkg->send_map_starts[par_fine_mat->commpkg->num_sends]];
        idx_t * recv_buf = new idx_t [par_fine_mat->offd.ncols];
        for (idx_t s = 0; s < par_fine_mat->commpkg->num_sends; s++) {
            idx_t       * buf    = send_buf                              + par_fine_mat->commpkg->send_map_starts[s];
            const idx_t * my_ids = par_fine_mat->commpkg->send_map_elmts + par_fine_mat->commpkg->send_map_starts[s];
            const idx_t num = par_fine_mat->commpkg->send_map_starts[s + 1] - par_fine_mat->commpkg->send_map_starts[s];
            for (idx_t k = 0; k < num; k++) {
                buf[k] = glb_map_f2c[my_ids[k]];
            }
            MPI_Isend(buf, num, mpi_idx_type, par_fine_mat->commpkg->send_pids[s], 114, comm, par_fine_mat->commpkg->send_reqs + s);
        }
        for (idx_t r = 0; r < par_fine_mat->commpkg->num_recvs; r++) {
            idx_t * buf = recv_buf + par_fine_mat->commpkg->recv_vec_starts[r];
            const idx_t num = par_fine_mat->commpkg->recv_vec_starts[r + 1] - par_fine_mat->commpkg->recv_vec_starts[r];
            MPI_Irecv(buf, num, mpi_idx_type, par_fine_mat->commpkg->recv_pids[r], 114, comm, par_fine_mat->commpkg->recv_reqs + r);
        }
        MPI_Waitall(par_fine_mat->commpkg->num_recvs, par_fine_mat->commpkg->recv_reqs, MPI_STATUSES_IGNORE);
    
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
            if (loc_map_f2c[loc_fi] == UNAGGREGATED) {// 该点仍未被聚合
                // 从所有邻居中选最大的进行聚合：注意包含落在别的进程上的邻居
                // 找出diag内最强耦合的邻居
                idx_t diag_candi = -1; setup_t diag_v = 0.0;// idx_t dc_ptr = -1;
                for (idx_t p = S_mat.row_ptr[loc_fi]; p < S_mat.row_ptr[loc_fi + 1]; p ++) {// 从大到小考察邻居
                    const idx_t loc_fj = S_mat.col_idx[p];
                    const idx_t glb_fj = trans_loc2glb(loc_fj);
                    if (loc_fj != loc_fi && // 该邻居不能是自己
                        iso_wells.find(glb_fj) == iso_wells.end() ) // 也不能是井
                    {
                        assert(glb_map_f2c[loc_fj] != -1);
                        diag_candi = glb_map_f2c[loc_fj];
                        diag_v     = S_mat.vals[p];
                        break;
                    }
                }
                // 找出offd内最强耦合的邻居
                const idx_t num_offd =  par_fine_mat->offd.row_ptr[loc_fi + 1] -
                                        par_fine_mat->offd.row_ptr[loc_fi    ]  ;
                setup_t offd_vbuf[ num_offd ];
                idx_t   offd_ibuf[ num_offd ];
                for (idx_t p = par_fine_mat->offd.row_ptr[loc_fi]; p < par_fine_mat->offd.row_ptr[loc_fi + 1]; p ++) {
                    const idx_t glb_fj = par_fine_mat->col_map_offd[ par_fine_mat->offd.col_idx[p] ];
                    if (iso_wells.find(glb_fj) == iso_wells.end()) {// 正常点计算范数
                        offd_vbuf[p - par_fine_mat->offd.row_ptr[loc_fi]] = calc_blk_norm(par_fine_mat->offd.vals + p * e_size);
                    } else {// 否则该细点为独立点，为避免与它聚合，联通值设为0
                        offd_vbuf[p - par_fine_mat->offd.row_ptr[loc_fi]] = 0.0;
                    }
                    offd_ibuf[p - par_fine_mat->offd.row_ptr[loc_fi]] = par_fine_mat->offd.col_idx[p];
                }
#ifdef PRES_COL_ORDER // 保证列号升序（以便coupled和naive在单进程时严格复现）
                sort_vals_with_indices<idx_t, setup_t>(num_offd, offd_vbuf, offd_ibuf);
#else // 快排不稳定排序
                my_qsort1<idx_t, setup_t, idx_t, -1>(offd_vbuf, offd_ibuf, 0, num_offd - 1);
#endif
                idx_t offd_candi = -1; setup_t offd_v = 0.0;// idx_t oc_ptr = -1;
                for (idx_t p = 0; p < num_offd; p ++) {// 从大到小考察邻居
                    const idx_t loc_fj = offd_ibuf[p];
                    const idx_t glb_cj = recv_buf[loc_fj];
                    if (glb_cj != -1)// 不能是邻居进程里尚未聚合的点
                    {
                        offd_candi = glb_cj;
                        offd_v     = offd_vbuf[p];
                        break;
                    }
                }
                idx_t final_candi = -1;// 最终确定的聚合邻居
                if (diag_candi != -1 && offd_candi != -1) {
                    if (diag_v > offd_v) final_candi = diag_candi;// 选值大的
                    else                 final_candi = offd_candi;
                }
                else { final_candi = MAX(diag_candi, offd_candi); }
                glb_map_f2c[loc_fi] = final_candi;
                if (final_candi == -1) {// check: 此时还尚未被聚合的，一定是只有一个邻居（即自己的）
                    assert(             S_mat.row_ptr[loc_fi + 1] -              S_mat.row_ptr[loc_fi] == 1);
                    assert(             S_mat.col_idx[S_mat.row_ptr[loc_fi]] == loc_fi);
                }
            }
        }
        MPI_Waitall(par_fine_mat->commpkg->num_sends, par_fine_mat->commpkg->send_reqs, MPI_STATUSES_IGNORE);
        delete [] recv_buf; delete [] send_buf;
    }// 此时在glb_map_f2c[]中仍为-1的点将不被粗化进入下一层粗网格

    {// 构造插值矩阵
        P_mat->diag.nrows = my_num_F; P_mat->offd.nrows = my_num_F;
        P_mat->diag.row_ptr = new idx_t [my_num_F + 1]; P_mat->diag.row_ptr[0] = 0;
        P_mat->offd.row_ptr = new idx_t [my_num_F + 1]; P_mat->offd.row_ptr[0] = 0;
        std::unordered_map<idx_t, std::pair<int, idx_t> > Poffd_g2l;// 记录插值矩阵的非对角部分的全局序到局部序的映射
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
            const idx_t glb_ci = glb_map_f2c[loc_fi];
            if (_cbeg <= glb_ci && glb_ci < _cend) {// 粗点属于本进程负责计算
                P_mat->diag.row_ptr[loc_fi + 1] = 1;
                P_mat->offd.row_ptr[loc_fi + 1] = 0;
            } else if (glb_ci >= 0) {// 粗点不属于本进程计算
                P_mat->diag.row_ptr[loc_fi + 1] = 0;
                P_mat->offd.row_ptr[loc_fi + 1] = 1;
                
                std::pair<int, idx_t> tmp(-1, -1);
                Poffd_g2l.emplace(glb_ci, tmp);
            } else { assert(glb_ci == -1);// 粗点不进入下一层
                P_mat->diag.row_ptr[loc_fi + 1] = 0;
                P_mat->offd.row_ptr[loc_fi + 1] = 0;
            }
        }
        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {// prefix sum
            P_mat->diag.row_ptr[loc_fi + 1] += P_mat->diag.row_ptr[loc_fi];
            P_mat->offd.row_ptr[loc_fi + 1] += P_mat->offd.row_ptr[loc_fi];
        }
        P_mat->prepare_commpkg(Poffd_g2l);

        // 确定列数和非零元数
        P_mat->diag.ncols = my_num_C;                      P_mat->offd.ncols = Poffd_g2l.size();
        P_mat->diag.nnz   = P_mat->diag.row_ptr[my_num_F]; P_mat->offd.nnz   = P_mat->offd.row_ptr[my_num_F];
        // 分配列序号和值的内存
        P_mat->diag.col_idx = new idx_t  [P_mat->diag.nnz]; P_mat->offd.col_idx = new idx_t  [P_mat->offd.nnz];
        P_mat->diag.vals    = new setup_t[P_mat->diag.nnz]; P_mat->offd.vals    = new setup_t[P_mat->offd.nnz];

        for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
            const idx_t glb_ci = glb_map_f2c[loc_fi];
            // const idx_t dcnt = 0, ocnt = 0;
            const idx_t doff = P_mat->diag.row_ptr[loc_fi],
                        ooff = P_mat->offd.row_ptr[loc_fi];
            if (_cbeg <= glb_ci && glb_ci < _cend) {// 粗点属于本进程负责计算
                P_mat->diag.col_idx[doff] = glb_ci - _cbeg;
                P_mat->diag.vals   [doff] = 1.0;
            } else if (glb_ci >= 0) {// 粗点不属于本进程计算
                assert(Poffd_g2l.find(glb_ci) != Poffd_g2l.end());
                P_mat->offd.col_idx[ooff] = Poffd_g2l[glb_ci].second;
                P_mat->offd.vals   [ooff] = 1.0;
            } else { assert(glb_ci == -1);// 粗点不进入下一层
                // nothing to do
            }
        }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Prlg diag %d %d %d offd %d %d %d\n", my_pid, P_mat->diag.nrows, P_mat->diag.ncols, P_mat->diag.nnz,
                    P_mat->offd.nrows, P_mat->offd.ncols, P_mat->offd.nnz);
                for (idx_t i = 0; i < P_mat->diag.nrows; i++) {
                    printf("  %d : ", P_mat->beg_row + i);
                    for (idx_t p = P_mat->diag.row_ptr[i]; p < P_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", P_mat->beg_col + P_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = P_mat->offd.row_ptr[i]; p < P_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", P_mat->col_map_offd[P_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }
    par_CSRMatrixTranspose(* P_mat, * R_mat);
    {// 需要修改限制矩阵的值
        for (idx_t i = 0; i < R_mat->diag.nrows; i ++) {
            const idx_t num = R_mat->diag.row_ptr[i + 1] - R_mat->diag.row_ptr[i]
                            + R_mat->offd.row_ptr[i + 1] - R_mat->offd.row_ptr[i];
            const setup_t v = 1.0 / (setup_t) num;
            for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p ++) {
                assert(R_mat->diag.vals[p] == 1.0);
                R_mat->diag.vals[p] = v;
            }
            for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p ++) {
                assert(R_mat->offd.vals[p] == 1.0);
                R_mat->offd.vals[p] = v;
            }
        }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d Rstr diag %d %d %d offd %d %d %d\n", my_pid, R_mat->diag.nrows, R_mat->diag.ncols, R_mat->diag.nnz,
                    R_mat->offd.nrows, R_mat->offd.ncols, R_mat->offd.nnz);
                for (idx_t i = 0; i < R_mat->diag.nrows; i++) {
                    printf("  %d : ", R_mat->beg_row + i);
                    for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->beg_col + R_mat->diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = R_mat->offd.row_ptr[i]; p < R_mat->offd.row_ptr[i + 1]; p++)
                        printf("%d ", R_mat->col_map_offd[R_mat->offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }
    delete [] glb_map_f2c;
#else // 最后的清扫只发生在进程内
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi++) {
        if (loc_map_f2c[loc_fi] == UNAGGREGATED) {// 该点仍未被聚合
            for (idx_t p = S_mat.row_ptr[loc_fi]; p < S_mat.row_ptr[loc_fi + 1]; p ++) {
                const idx_t loc_fj = S_mat.col_idx[p];
                if (loc_fj != loc_fi &&
                    iso_wells.find(trans_loc2glb(loc_fj)) == iso_wells.end()) {
                    assert(loc_map_f2c[loc_fj] != UNAGGREGATED);
                    // printf("sweep F %d with F %d into C %d\n", loc_fi, loc_fj, loc_map_f2c[loc_fj]);
                    loc_map_f2c[loc_fi] = loc_map_f2c[loc_fj];// 归入该邻居
                    break;
                } else assert(S_mat.vals[p] != 0.0);
            }
        }
        // 此时还尚未被聚合的，一定是只有一个邻居（即自己的）
        if (loc_map_f2c[loc_fi] == UNAGGREGATED) {
            assert(S_mat.row_ptr[loc_fi + 1] - S_mat.row_ptr[loc_fi] == 1);
            assert(S_mat.col_idx[S_mat.row_ptr[loc_fi]] == loc_fi);
        }
    }
    // 构造插值矩阵
    P_mat->diag.nrows = my_num_F; P_mat->offd.nrows = my_num_F;// 非对角部分行数照常
    P_mat->diag.ncols = my_num_C; P_mat->offd.ncols = 0;// 非对角部分列数为0
    P_mat->diag.row_ptr = new idx_t [my_num_F + 1]; P_mat->diag.row_ptr[0] = 0;
    P_mat->offd.row_ptr = new idx_t [my_num_F + 1]; P_mat->offd.row_ptr[0] = 0;
#if 1 // 多线程
    idx_t twspace [omp_get_max_threads()];
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        int  nt = omp_get_num_threads();
        idx_t tbeg = -1, tend = -1;
        thread_load_balance(my_num_F, nt, tid, tbeg, tend);

        idx_t tcnt = 0;
        for (idx_t loc_fi = tbeg; loc_fi < tend; loc_fi ++) {
            idx_t loc_ci = loc_map_f2c[loc_fi];
            if (loc_ci == UNAGGREGATED) {// 此时还尚未被聚合的，一定是只有一个邻居（即自己的），它们不需要进入到粗网格，直接在细层求解即可（即独立点）
                assert(S_mat.row_ptr[loc_fi + 1] - S_mat.row_ptr[loc_fi] == 1);
                assert(S_mat.col_idx[S_mat.row_ptr[loc_fi]] == loc_fi);
                P_mat->diag.row_ptr[loc_fi + 1] = tcnt;// 插值矩阵里该行为空
            } else {
                tcnt ++;
                P_mat->diag.row_ptr[loc_fi + 1] = tcnt;// plain
            }
            P_mat->offd.row_ptr[loc_fi + 1] = 0;
        }
        twspace[tid] = tcnt;

        #pragma omp barrier
        if (tid > 0) {
            idx_t prev_t_sum = 0;// prefix sum
            for (int prev_t = 0; prev_t < tid; prev_t ++) prev_t_sum += twspace[prev_t];
            for (idx_t loc_fi = tbeg; loc_fi < tend; loc_fi ++)
                P_mat->diag.row_ptr[loc_fi + 1] += prev_t_sum;
        } else {// tid == 0
            idx_t tot_t_sum = 0;// all threads sum
            for (int t = 0; t < nt; t ++) tot_t_sum += twspace[t];
            P_mat->diag.nnz = tot_t_sum;
            P_mat->diag.col_idx = new idx_t  [P_mat->diag.nnz];
            P_mat->diag.vals    = new setup_t[P_mat->diag.nnz];
        }
        
        #pragma omp barrier
        
        for (idx_t loc_fi = tbeg; loc_fi < tend; loc_fi ++) {
            idx_t loc_ci = loc_map_f2c[loc_fi];
            if (loc_ci != UNAGGREGATED) {
                const idx_t off = P_mat->diag.row_ptr[loc_fi];
                P_mat->diag.col_idx[off] = loc_ci;
                P_mat->diag.vals   [off] = 1.0;// plain
            }
        }
    }
#else // 串行
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
        idx_t loc_ci = loc_map_f2c[loc_fi];
        if (loc_ci == UNAGGREGATED) {// 此时还尚未被聚合的，一定是只有一个邻居（即自己的），它们不需要进入到粗网格，直接在细层求解即可（即独立点）
            assert(S_mat.row_ptr[loc_fi + 1] - S_mat.row_ptr[loc_fi] == 1);
            assert(S_mat.col_idx[S_mat.row_ptr[loc_fi]] == loc_fi);
            P_mat->diag.row_ptr[loc_fi + 1] = 0;// 插值矩阵里该行为空
        } else {
            P_mat->diag.row_ptr[loc_fi + 1] = 1;// plain
        }
        P_mat->offd.row_ptr[loc_fi + 1] = 0;
    }

    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
        P_mat->diag.row_ptr[loc_fi + 1] += P_mat->diag.row_ptr[loc_fi];
    }
    P_mat->diag.nnz = P_mat->diag.row_ptr[my_num_F];
    P_mat->diag.col_idx = new idx_t  [P_mat->diag.nnz];
    P_mat->diag.vals    = new setup_t[P_mat->diag.nnz];
    for (idx_t loc_fi = 0; loc_fi < my_num_F; loc_fi ++) {
        idx_t loc_ci = loc_map_f2c[loc_fi];
        if (loc_ci != UNAGGREGATED) {
            const idx_t off = P_mat->diag.row_ptr[loc_fi];
            P_mat->diag.col_idx[off] = loc_ci;
            P_mat->diag.vals   [off] = 1.0;// plain
        }
    }
#endif
    // 转置构造限制矩阵
    CSRMatrixTranspose(P_mat->diag, R_mat->diag);
    R_mat->offd.nrows = R_mat->diag.nrows;// 非对角部分行数照常
    R_mat->offd.ncols = 0;// 非对角部分列数为0
    R_mat->offd.row_ptr = new idx_t [R_mat->offd.nrows + 1]; R_mat->offd.row_ptr[0] = 0;
    // 需要修改限制矩阵的值
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < R_mat->diag.nrows; i++) {
        const idx_t num = R_mat->diag.row_ptr[i + 1] - R_mat->diag.row_ptr[i]; assert(num > 0);
        const setup_t v = 1.0 / (setup_t) num;
        for (idx_t p = R_mat->diag.row_ptr[i]; p < R_mat->diag.row_ptr[i + 1]; p++) {
            assert(R_mat->diag.vals[p] == 1.0);
            R_mat->diag.vals[p] = v;
        }
        R_mat->offd.row_ptr[i + 1] = 0;
    }
#endif // 最后的清扫是否可跨进程
// #ifdef DEBUG
    const idx_t my_num_I = my_num_F - P_mat->diag.nnz - P_mat->offd.nnz;// 独立点个数
    idx_t tot_num_I = 0;
    MPI_Reduce(& my_num_I, & tot_num_I, 1, mpi_idx_type, MPI_SUM, 0, comm);
    if (my_pid == 0) printf("Glb #I %d\n", tot_num_I);
// #endif
    delete [] loc_map_f2c;
    return C_rows_partits;
}

template int* naive_pairwise(const par_CSRMatrix_I32F64F64_DOF1 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);
template int* naive_pairwise(const par_CSRMatrix_I32F64F64_DOF2 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);
template int* naive_pairwise(const par_CSRMatrix_I32F64F64_DOF3 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);
template int* naive_pairwise(const par_CSRMatrix_I32F64F64_DOF4 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &);

template int* naive_pairwise(const par_CSRMatrix_I32F32F32_DOF1 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);
template int* naive_pairwise(const par_CSRMatrix_I32F32F32_DOF2 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);
template int* naive_pairwise(const par_CSRMatrix_I32F32F32_DOF3 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);
template int* naive_pairwise(const par_CSRMatrix_I32F32F32_DOF4 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &);

